#include "Waves.h"

#include <algorithm>

#include "Vertex.h"

namespace D3D10Utils
{
	Waves::Waves(const unsigned __int32 rows, const unsigned __int32 columns)
		: mDevice(0)
		, mVertexBuffer(0)
		, mIndexBuffer(0)
		, mPreviousSolution(0)
		, mCurrentSolution(0)
		, mNormals(0)
		, mNumVertices(rows * columns)
		, mNumFaces((rows - 1) * (columns - 1) * 2)
		, mRows(rows)
		, mColumns(columns)
		, mK1(0.0f)
		, mK2(0.0f)
		, mK3(0.0f)
		, mTimeStep(0.0f)
		, mSpatialStep(0.0f)
	{

	}

	Waves::~Waves()
	{
		delete[] mPreviousSolution;
		delete[] mCurrentSolution;
		delete[] mNormals;

		ReleaseCOM(mVertexBuffer);
		ReleaseCOM(mIndexBuffer);
	}

	void Waves::init(ID3D10Device * const device, const float cellSpacing, const float deltaTime, 
		const float speed, const float damping)
	{
		if (!device)
		{
			MessageBox(0, L"NULL pointer", 0, 0);
			PostQuitMessage(0);
		}

		mDevice = device;

		mTimeStep = deltaTime;
		mSpatialStep = cellSpacing;

		const float d = damping * deltaTime + 2.0f;
		const float e = (speed * speed) * (deltaTime * deltaTime) / (cellSpacing * cellSpacing);
		mK1 = (damping * deltaTime - 2.0f) / d;
		mK2 = (4.0f - 8.0f * e) / d;
		mK3 = (2.0f * e) / d;

		mPreviousSolution = new D3DXVECTOR3[mNumVertices];
		mCurrentSolution = new D3DXVECTOR3[mNumVertices];	
		mNormals = new D3DXVECTOR3[mNumVertices];		

		initVertexBuffer();
		initIndexBuffer();
	}

	void Waves::update(const float deltaTime)
	{
		static float t = 0;

		// Accumulate time.
		t += deltaTime;

		// Only update the simulation at the specified time step.
		if (t >= mTimeStep)
		{
			// Only update interior points; we use zero boundary conditions.
			for(unsigned __int32 i = 1; i < mRows - 1; ++i)
			{
				for(unsigned __int32 j = 1; j < mColumns - 1; ++j)
				{
					// After this update we will be discarding the old previous
					// buffer, so overwrite that buffer with the new update.
					// Note how we can do this inplace (read/write to same element) 
					// because we won't need prev_ij again and the assignment happens last.

					// Note j indexes x and i indexes z: h(x_j, z_i, t_k)
					// Moreover, our +z axis goes "down"; this is just to 
					// keep consistent with our row indices going down.
					mPreviousSolution[i * mRows + j].y = 
						mK1*mPreviousSolution[i *mRows + j].y +
						mK2*mCurrentSolution[i * mRows + j].y +
						mK3*(mCurrentSolution[(i + 1) * mRows + j].y + 
						mCurrentSolution[(i - 1) * mRows + j].y + 
						mCurrentSolution[i * mRows + j + 1].y + 
						mCurrentSolution[i * mRows + j - 1].y);
				}
			}

			// We just overwrote the previous buffer with the new data, so
			// this data needs to become the current solution and the old
			// current solution becomes the new previous solution.
			std::swap(mPreviousSolution, mCurrentSolution);

			t = 0.0f; // reset time

			// Compute normals using finite difference scheme.
			for(unsigned __int32 i = 1; i < mRows - 1; ++i)
			{
				for(unsigned __int32 j = 1; j < mColumns - 1; ++j)
				{
					const float l = mCurrentSolution[i *mColumns + j - 1].y;
					const float r = mCurrentSolution[i * mColumns + j + 1].y;
					const float t = mCurrentSolution[(i - 1) * mColumns + j - 1].y;
					const float b = mCurrentSolution[(i + 1) * mColumns + j + 1].y;
					mNormals[i * mColumns + j].x = -r + l;
					mNormals[i * mColumns + j].y = 2.0f * mSpatialStep;
					mNormals[i * mColumns + j].z = b - t;

					D3DXVec3Normalize(&mNormals[i * mColumns + j], &mNormals[i * mColumns + j]);
				}
			}

			// Update the vertex buffer with the new solution.
			Vertex* v = 0;
			ErrorHandler(mVertexBuffer->Map(D3D10_MAP_WRITE_DISCARD, 0, 
				reinterpret_cast<void**> (&v)));

			// Compute the delta between texels according to
			// waves grid dimensions.
			const float du = 1.0f / (mColumns - 1);
			const float dv = 1.0f / (mRows - 1);

			for(unsigned __int32 i = 0; i < mRows; ++i)
			{
				for(unsigned __int32 j = 0; j < mColumns; ++j)
				{
					const unsigned __int32 vertexIndex = i * mColumns + j;
					v[vertexIndex].mPosition    = mCurrentSolution[vertexIndex];
					v[vertexIndex].mNormal = mNormals[vertexIndex];
					v[vertexIndex].mTexCoord.x = j * du;
					v[vertexIndex].mTexCoord.y = i * dv;
				}
			}
			
			mVertexBuffer->Unmap();
		}
	}

	void Waves::disturb(const unsigned __int32 iRow, const unsigned __int32 jColumn, const float magnitude)
	{
		// Don't disturb boundaries.
		assert(iRow > 0 && iRow < mRows - 1);
		assert(jColumn > 0 && jColumn < mColumns - 1);

		const float halfMag = 0.5f * magnitude;

		// Disturb the ijth vertex height and its neighbors.
		mCurrentSolution[iRow * mRows + jColumn].y += magnitude;
		mCurrentSolution[iRow * mRows + jColumn + 1].y += halfMag;
		mCurrentSolution[iRow * mRows + jColumn - 1].y += halfMag;
		mCurrentSolution[(iRow + 1) * mRows + jColumn].y += halfMag;
		mCurrentSolution[(iRow - 1) * mRows + jColumn].y += halfMag;
	}

	void Waves::draw()
	{
		const unsigned __int32 stride = sizeof(Vertex);
		const unsigned __int32 offset = 0;
		mDevice->IASetVertexBuffers(0, 1, &mVertexBuffer, &stride, &offset);
		mDevice->IASetIndexBuffer(mIndexBuffer, DXGI_FORMAT_R32_UINT, 0);

		mDevice->DrawIndexed(mNumFaces * 3, 0, 0);
	}

	Vertex* Waves::createVertices()
	{
		// Compute the half width and height of the grid.		
		// Our grid will be at xz plane, that the reason why we ask for cell
		// spacing in that coordinates.
		const float halfWidth = (mColumns - 1) * mSpatialStep * 0.5f;
		const float halfHeight = (mRows - 1) * mSpatialStep * 0.5f;

		// Generate grid vertices in system memory.
		// Vij = (-0.5 * width + i * cellSpacingX, 0.5 * height - j * cellSpacingZ)
		Vertex *vertices = new Vertex[mNumVertices];
		for(unsigned __int32 j = 0; j < mColumns; ++j)
		{
			const float x = -halfWidth + j * mSpatialStep;

			for(UINT i = 0; i < mRows; ++i)
			{
				const float z = halfHeight - i * mSpatialStep;

				// Update the vertex
				const unsigned __int32 vertexIndex = i * mRows + j;
				mPreviousSolution[vertexIndex] = D3DXVECTOR3(x, 0.0f, z);
				mCurrentSolution[vertexIndex] = D3DXVECTOR3(x, 0.0f, z);
				mNormals[vertexIndex] = D3DXVECTOR3(0.0f, 1.0f, 0.0f);
			}
		}

		return vertices;
	}

	unsigned __int32* Waves::createIndices()
	{
		unsigned __int32 *indices = new unsigned __int32[mNumFaces * 3];
		unsigned __int32 quadOffset = 0;
		for(unsigned __int32 j = 0; j < mColumns - 1; ++j)
		{
			for(unsigned __int32 i = 0; i < mRows - 1; ++i)
			{
				// Define the needed vertices.
				const unsigned __int32 upperLeft = i * mRows + j;
				const unsigned __int32 upperRight = i * mRows + j + 1;
				const unsigned __int32 bottomLeft = (i + 1) * mRows + j;
				const unsigned __int32 bottomRight = (i + 1) * mRows + j + 1;

				indices[quadOffset] = upperLeft;
				indices[quadOffset + 1] = upperRight;
				indices[quadOffset + 2] = bottomLeft;

				indices[quadOffset + 3] = bottomLeft;
				indices[quadOffset + 4] = upperRight;
				indices[quadOffset + 5] = bottomRight;

				// Next quad
				quadOffset += 6; 
			}
		}

		return indices;
	}

	void Waves::initVertexBuffer()
	{
		Vertex *vertices = createVertices();

		D3D10_BUFFER_DESC vertexBufferDescription;
		vertexBufferDescription.ByteWidth = sizeof(Vertex) * mNumVertices;
		vertexBufferDescription.Usage = D3D10_USAGE_DYNAMIC; 
		vertexBufferDescription.BindFlags = D3D10_BIND_VERTEX_BUFFER; 
		vertexBufferDescription.CPUAccessFlags = D3D10_CPU_ACCESS_WRITE; 
		vertexBufferDescription.MiscFlags = 0; 	

		D3D10_SUBRESOURCE_DATA vertexBufferData;
		vertexBufferData.pSysMem = vertices;

		HRESULT hr = mDevice->CreateBuffer(&vertexBufferDescription, &vertexBufferData, &mVertexBuffer);
		ErrorHandler(hr);

		delete[] vertices;
	}

	void Waves::initIndexBuffer()
	{
		unsigned __int32 *indices = createIndices();

		D3D10_BUFFER_DESC indexBufferDescription;
		indexBufferDescription.ByteWidth = sizeof(unsigned __int32) * mNumFaces * 3;
		indexBufferDescription.Usage = D3D10_USAGE_IMMUTABLE;
		indexBufferDescription.BindFlags = D3D10_BIND_INDEX_BUFFER;
		indexBufferDescription.CPUAccessFlags = 0; 
		indexBufferDescription.MiscFlags = 0;

		D3D10_SUBRESOURCE_DATA indexBufferData;
		indexBufferData.pSysMem = indices;

		HRESULT hr = mDevice->CreateBuffer(&indexBufferDescription, &indexBufferData, &mIndexBuffer);
		ErrorHandler(hr);

		delete[] indices;
	}
}